#!/usr/bin/env bash

# 以下这些镜像也是存在的，默认不下载，若要下载，则维护到base_image_list中
# nginx_1.26.3
# node_18-alpine
# node_20-alpine
# eclipse-temurin_8u432-b06-jre
# eclipse-temurin_8u442-b06-jdk
# eclipse-temurin_21-jdk-alpine
# dragonwell_8.24.23-standard-ga-ubuntu
# libreoffice-tesseract-ocr5-bookworm_dragonwell11-slim
base_image_list=(
    nginx_1.28.0
    nginx_stable-alpine
    node_14-alpine
    eclipse-temurin_21-jre-alpine
    dragonwell_21.0.5.0.5-bookworm
    python_3.11-slim-bullseye-nltk
    buildkit_buildx-stable-1
    libreoffice-tesseract-ocr5-bookworm_temurin21-jre
    qdrant_v1.11.0
)

# 原镜像地址（默认值，不修改）
declare -r rddhub_registry="rddhub.changhong.com/library"
# 私有仓库镜像地址（根据实际情况修改，需要先登录）
declare -r private_registry="registry.cn-east-3.myhuaweicloud.com/library"
# 华为obs基地址
declare -r obs_base_url="https://brianhsiung.obs.cn-east-3.myhuaweicloud.com/images/base"

DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$DIR" || exit 1

##############################################
# 日志函数
##############################################
log_info()    { echo -e "\033[34m[INFO] $*\033[0m"; }
log_warn()    { echo -e "\033[33m[WARNING] $*\033[0m"; }
log_success() { echo -e "\033[32m[SUCCESS] $*\033[0m"; }
log_error()   { echo -e "\033[31m[ERROR] $*\033[0m"; exit 1; }

##############################################
# 检查依赖
##############################################
check_dependency() {

    local -a required_cmds=("wget" "docker")
    local missing=()

    log_info "检查依赖命令: ${required_cmds[*]}"
    for cmd in "${required_cmds[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    ((${#missing[@]})) && log_error "缺失依赖命令: ${missing[*]}"
}

##############################################
# 处理单个镜像
##############################################
process_image() {
    local image_name="$1"
    local image_tag="$2"
    local huawei_obs_url="$3"
    local private_image="$4"

    local image_file="${image_name}_${image_tag}.tar.gz"
    local image_url="${huawei_obs_url}/${image_file}"

    if [ ! -f "$image_file" ]; then
        log_info "下载镜像: $image_file"
        wget -q "$image_url" || log_error "下载失败: $image_url"
        #curl -o "$image_file" "$image_url" || log_error "下载失败: $image_url"
        log_success "下载成功: $image_file"
    else
        log_info "镜像已存在: $image_file"
    fi

    log_info "加载镜像: $image_file"
    sudo docker load -i "$image_file" || log_error "加载镜像失败: $image_file"
    sudo docker tag "$rddhub_registry/${image_name}:${image_tag}" "$private_image"
    sudo docker push "$private_image" || log_error "推送镜像失败: $private_image"
}

##############################################
# 获取单CPU基础镜像
##############################################
get_single_base_image() {

    arch="$(uname -m)"
    case "$arch" in
        aarch64|arm64)
            huawei_obs_url="${obs_base_url}/aarch64"
        ;;
        amd64|x86_64)
            huawei_obs_url="${obs_base_url}/x64"
        ;;
        *)
            log_error "不支持的CPU架构: $arch"
        ;;
    esac

    for image in "${base_image_list[@]}"; do
        # 从左到右匹配第一个下划线前的所有字符
        local image_name="${image%%_*}"
        # 从左到右匹配第一个下划线之后的所有字符
        local image_tag="${image#*_}"
        # 私有镜像仓库地址
        local private_image="${private_registry}/${image_name}:${image_tag}"
        
        process_image "$image_name" "$image_tag" "$huawei_obs_url" "$private_image"
    done
}

##############################################
# 获取x86_64和arm64基础镜像
##############################################
get_muliarch_base_image() {
    [ -d "$DIR/arm64" ] || mkdir -p "$DIR/arm64"
    [ -d "$DIR/amd64" ] || mkdir -p "$DIR/amd64"

    declare -A arch_urls=(
        ["amd64"]="${obs_base_url}/x64"
        ["arm64"]="${obs_base_url}/aarch64"
    )

    for arch in "${!arch_urls[@]}"; do
        cd "$DIR/$arch" || exit 1
        for image in "${base_image_list[@]}"; do
            # 从左到右匹配第一个下划线前的所有字符
            local image_name="${image%%_*}"
            # 从左到右匹配第一个下划线之后的所有字符
            local image_tag="${image#*_}"
            # 私有镜像仓库地址
            local private_image="${private_registry}/${image_name}:${image_tag}-${arch}"

            process_image "$image_name" "$image_tag" "${arch_urls[$arch]}" "$private_image"
        done
    done

    for image in "${base_image_list[@]}"; do
        # 从左到右匹配第一个下划线前的所有字符
        local image_name="${image%%_*}"
        # 从左到右匹配第一个下划线之后的所有字符
        local image_tag="${image#*_}"
        # 私有镜像仓库地址
        local private_arm_image="${private_registry}/${image_name}:${image_tag}-arm64"
        local private_amd_image="${private_registry}/${image_name}:${image_tag}-amd64"

        sudo docker manifest create "$private_registry/${image_name}:${image_tag}" \
            "$private_arm_image" "$private_amd_image" || log_error "创建镜像清单失败: $private_registry/$image_name:$image_tag"
        sudo docker manifest push "$private_registry/${image_name}:${image_tag}" || log_error "推送镜像清单失败: $private_registry/$image_name:$image_tag"
        #sudo docker manifest inspect "$private_registry/${image_name}:${image_tag}" || log_error "检查镜像清单失败: $private_registry/$image_name:$image_tag"
        sudo docker manifest rm "$private_registry/${image_name}:${image_tag}" || log_error "删除镜像清单失败: $private_registry/$image_name:$image_tag"
    done

    log_info "创建用于构建x86_64和arm64镜像的容器"
    sudo docker buildx create \
        --name=multi-platform-builder \
        --platform=linux/amd64,linux/arm64 \
        --driver=docker-container \
        --driver-opt image="$private_registry"/buildkit:buildx-stable-1 \
        --use \
        --bootstrap
}

main() {
    
    check_dependency

    while :
    do
        read -p "是否下载单一CPU架构镜像 ( Y/y 表示下载单一CPU架构镜像, N/n 表示下载多架构镜像): " -r option
        case "${option}" in
            Y|y)
                get_single_base_image && log_success "下载成功" && exit 0
            ;;
            N|n)
                get_muliarch_base_image && log_success "下载成功" && exit 0
            ;;
            *)
                log_warn "无效输入，请输入 Y(y) 或 N(n)"
            ;;
        esac
    done
}

main "$@"